from typing import Tuple

import numpy as np
import pandas as pd
import pandas_ta as ta

from .indicator import Indicator


class SmaIndicator(Indicator):
    """计算简单移动平均 (SMA)。"""

    def __init__(self, source_column: Tuple[str, str], period: int, name='sma'):
        if period < 1:
            raise ValueError("周期必须大于0。")
        super().__init__(source_column=source_column, period=period, name=name)

    def compute(self, df: pd.DataFrame) -> pd.Series:
        sma = ta.sma(df[self.source_column], length=self.params['period'])
        print(sma)
        return sma


class StdIndicator(Indicator):
    """计算滚动标准差。"""

    def __init__(self, source_column: Tuple[str, str], period: int):
        if period < 1:
            raise ValueError("周期必须大于0。")
        super().__init__(source_column=source_column, period=period)

    def compute(self, df: pd.DataFrame) -> pd.Series:
        return df[self.source_column].rolling(window=self.params['period']).std()


class ZScoreIndicator(Indicator):
    """
    计算Z-Score。
    这是一个组合指标，它内部依赖于SMA和STD的计算逻辑。
    """

    def __init__(self, source_column: Tuple[str, str], period: int):
        if period < 1:
            raise ValueError("周期必须大于0。")
        super().__init__(source_column=source_column, period=period)

    def compute(self, df: pd.DataFrame) -> pd.Series:
        period = self.params['period']
        data_series = df[self.source_column]
        # 直接在输入序列上进行滚动计算，这是最高效的方式
        sma = data_series.rolling(window=period).mean()
        std = data_series.rolling(window=period).std()

        # 核心计算: (value - mean) / std_dev
        z_score = (data_series - sma) / std

        # 清理因std为0导致的无穷大值和计算初期的NaN值
        z_score.replace([np.inf, -np.inf], 0, inplace=True)
        z_score.fillna(0, inplace=True)

        return z_score


class ExpressionIndicator(Indicator):
    """
    一个特殊的指标，它的计算逻辑由一个字符串表达式定义。
    这允许用户在运行时创建临时的、复合的指标。
    """

    def __init__(self, source_column: Tuple[str, str], expression: str):
        """
        Args:
            expression (str): 一个有效的Pandas表达式，例如 "close - open"。
        """
        super().__init__(source_column=source_column, expression=expression)

    def compute(self, df: pd.DataFrame) -> pd.Series:
        """
        使用DataFrame的eval()方法来执行表达式。

        注意：这个指标的compute方法接收的是整个DataFrame，而不是单个Series，
        因为它需要访问表达式中可能存在的多个列。
        """
        try:
            # 我们将compute的签名推广为接收DataFrame
            return df.eval(self.params['expression'])
        except Exception as e:
            raise ValueError(f"执行表达式 '{self.params['expression']}' 失败: {e}")

    # 我们需要一种方式来命名这个临时指标的输出列
    @property
    def name(self) -> str:
        # 使用表达式本身作为名称的一部分，并进行清理
        safe_expr_name = self.params['expression'].replace(' ', '').replace('(', '').replace(')', '').replace('/',
                                                                                                              '_div_')
        return f"EXPR({safe_expr_name})"


class ATRIndicator(Indicator):
    """
    计算平均真实波幅 (Average True Range, ATR)。

    ATR是一个衡量市场波动性的技术指标。它需要高、低、收盘价作为输入。
    这个实现类将演示如何处理需要多个输入源列的指标。
    """

    def __init__(self, timeframe: str, period: int):
        """
        初始化ATR指标。

        Args:
            timeframe (str):
                要计算ATR的时间周期, e.g., "1h", "5m"。
                框架将使用这个字符串来定位 '1h_high', '1h_low', '1h_close' 等列。
            period (int): ATR的计算周期, e.g., 14。
        """
        if period < 1:
            raise ValueError("周期必须大于0。")

        # 对于需要多列的指标，我们将source_column设为None，
        # 因为它的输入源不是单一的一列。
        # 我们将timeframe和period作为参数保存。
        super().__init__(source_column=None, timeframe=timeframe, period=period)

    def _generate_column_name(self) -> Tuple[str, str]:
        """
        生成唯一的输出列名。
        例如: 对于1小时周期的14周期ATR，返回 ('1h', 'atr_14')。
        """
        timeframe = self.params['timeframe']
        period = self.params['period']

        # 我们将ATR指标的值与它所计算的时间周期关联在一起
        return timeframe, f"atr_{period}"

    def compute(self, df: pd.DataFrame) -> pd.Series:
        """
        使用pandas-ta库计算ATR值。
        """
        timeframe = self.params['timeframe']
        period = self.params['period']

        # 1. 根据timeframe动态构建所需的高、低、收盘价的列名（元组）
        high_col = (timeframe, 'high')
        low_col = (timeframe, 'low')
        close_col = (timeframe, 'close')

        # 2. 检查所需的所有列是否存在于DataFrame中
        required_cols = [high_col, low_col, close_col]
        for col in required_cols:
            if col not in df.columns:
                raise ValueError(f"ATR计算失败：所需的输入列 '{col}' 在DataFrame中不存在。")

        # 3. 从DataFrame中提取所需的数据序列
        highs = df[high_col]
        lows = df[low_col]
        closes = df[close_col]

        # 4. 调用 pandas-ta 的 atr 方法进行计算
        #    .squeeze() 方法可以将单列的DataFrame转换为Series，如果需要的话
        atr_series = ta.atr(high=highs, low=lows, close=closes, length=period)

        return atr_series